from typing import Dict, List
import random
from utils.Seed import Seed
from schedule.PowerSchedule import PowerSchedule

class CutoffExponentialSchedule(PowerSchedule):
    def __init__(self, beta: float = 2.0, M: float = 100.0) -> None:
        """
        截断指数调度器
        :param beta: 衰减因子 (beta > 1)
        :param M: 单一路径的最大能量限制
        """
        super().__init__()
        self.beta = beta
        self.M = M
        self.path_stats: Dict[int, dict] = {}  # 路径统计 {path_id: {'freq': int, 'selections': int}}
        self.global_mu = 0  # 全局平均测试频率

    def bind_frequency(self, freq_dict: Dict[int, int]) -> None:
        """绑定外部频率字典（兼容PathPowerSchedule接口）"""
        for path_id, freq in freq_dict.items():
            if path_id not in self.path_stats:
                self.path_stats[path_id] = {'freq': 0, 'selections': 0}
            self.path_stats[path_id]['freq'] = freq
        self._update_global_mu()

    def update_stats(self, path_id: int) -> None:
        """更新路径统计（测试执行后调用）"""
        if path_id not in self.path_stats:
            self.path_stats[path_id] = {'freq': 0, 'selections': 0}
        self.path_stats[path_id]['freq'] += 1
        self._update_global_mu()

    def mark_seed(self, path_id: int) -> None:
        """标记路径被选为种子（发现新路径时调用）"""
        if path_id in self.path_stats:
            self.path_stats[path_id]['selections'] += 1

    def _update_global_mu(self) -> None:
        """计算全局平均测试频率"""
        if not self.path_stats:
            self.global_mu = 0
            return
        self.global_mu = sum(stats['freq'] for stats in self.path_stats.values()) / len(self.path_stats)

    def assign_energy(self, population: List[Seed]) -> None:
        """核心能量分配逻辑"""
        if not self.path_stats:
            # 初始状态：均匀分配
            for seed in population:
                seed.energy = 1
            return

        for seed in population:
            path_id = hash(frozenset(seed.coverage))
            stats = self.path_stats.get(path_id, {'freq': 0, 'selections': 0})
            
            if stats['freq'] > self.global_mu:  # 高频路径
                seed.energy = 0
            else:  # 低频路径
                energy = (1 / self.beta) * (2 ** stats['selections'])
                seed.energy = min(energy, self.M)